#include "KCF.h"
#include <cmath>

KCFTracker::KCFTracker()
    : _cell_size(1), _padding(2.5), _lambda(1e-4), _output_sigma_factor(0.1) {}

void KCFTracker::init(const Eigen::Vector4f& roi, Image* image) {
    _tmpl = getFeatures(image, true);
    _prob = createGaussianPeak(_tmpl.rows(), _tmpl.cols());
    _alphaf = Eigen::MatrixXf::Zero(_tmpl.rows(), _tmpl.cols());
}

Eigen::Vector4f KCFTracker::update(Image* image) {
    Eigen::MatrixXf features = getFeatures(image, false);

    Eigen::MatrixXf kf = gaussianCorrelation(features, _tmpl);

    Eigen::MatrixXf response = Eigen::MatrixXf::Zero(_alphaf.rows(), _alphaf.cols());;
    Eigen::MatrixXcd kf_complex = _alphaf.cast<std::complex<double>>();

    Eigen::FFT<double> fft;

    // 执行二维傅里叶变换
    for (int i = 0; i < kf_complex.rows(); ++i) {
        Eigen::VectorXcd row_fft;
        fft.fwd(row_fft, kf_complex.row(i));
        response.row(i) = row_fft;
    }

    for (int i = 0; i < response.cols(); ++i) {
        Eigen::VectorXcd col_fft;
        fft.fwd(col_fft, response.col(i));
        response.col(i) = col_fft;
    }

    Eigen::MatrixXf::Index maxRow, maxCol;
    response.maxCoeff(&maxRow, &maxCol);

    Eigen::Vector4f new_roi;
    new_roi << maxCol - _tmpl_sz.x() / 2, maxRow - _tmpl_sz.y() / 2, _tmpl_sz.x(), _tmpl_sz.y();

    Eigen::MatrixXf new_features = getFeatures(image, false);
    Eigen::MatrixXf new_kf = gaussianCorrelation(new_features, new_features);
    _alphaf = _prob.cwiseProduct(new_kf + Eigen::MatrixXf::Constant(_tmpl.rows(), _tmpl.cols(), _lambda)).cwiseQuotient(new_kf + Eigen::MatrixXf::Constant(_tmpl.rows(), _tmpl.cols(), _lambda));
    _tmpl = new_features;

    return new_roi;
}

Eigen::MatrixXf KCFTracker::gaussianCorrelation(const Eigen::MatrixXf& x1, const Eigen::MatrixXf& x2) {
    
    Eigen::MatrixXcd x1_complex = x1.cast<std::complex<double>>();
    Eigen::MatrixXcd c(x1_complex.rows(), x1_complex.cols());

    Eigen::FFT<double> fft;

    // 执行二维傅里叶变换
    for (int i = 0; i < x1_complex.rows(); ++i) {
        Eigen::VectorXcd row_fft;
        fft.fwd(row_fft, x1_complex.row(i));
        c.row(i) = row_fft;
    }

    for (int i = 0; i < c.cols(); ++i) {
        Eigen::VectorXcd col_fft;
        fft.fwd(col_fft, c.col(i));
        c.col(i) = col_fft;
    }

    float d = (x1.array().square() + x2.array().square()).sum();
    Eigen::MatrixXf d_matrix = Eigen::MatrixXf::Constant(c.rows(), c.cols(), d);
    return (-1.0 / (_cell_size * _cell_size) * (d_matrix - 2.0 * c).cwiseMax(0)).array().exp();
}

Eigen::MatrixXf KCFTracker::createGaussianPeak(int sizey, int sizex) {
    Eigen::MatrixXf res(sizey, sizex);
    int syh = sizey / 2;
    int sxh = sizex / 2;

    float output_sigma = std::sqrt(static_cast<float>(sizex * sizey)) / _padding * _output_sigma_factor;
    float mult = -0.5 / (output_sigma * output_sigma);

    for (int i = 0; i < sizey; i++)
        for (int j = 0; j < sizex; j++) {
            int ih = i - syh;
            int jh = j - sxh;
            res(i, j) = std::exp(mult * static_cast<float>(ih * ih + jh * jh));
        }

    return res;
}

Eigen::MatrixXf KCFTracker::getFeatures(Image* image, bool inithann) {
    int width = image->Myformat.width;
    int height = image->Myformat.height;

    Eigen::MatrixXf features(height, width);
    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            double gray = (image->at(y, x, 1) + image->at(y, x, 2) + image->at(y, x, 3))/3;
            features(y, x) = static_cast<double>(gray) / 255.0f;
        }
    }

    if (inithann) {
        _hann = Eigen::MatrixXf(height, width);
        for (int i = 0; i < height; ++i)
            for (int j = 0; j < width; ++j)
                _hann(i, j) = 0.5 * (1 - std::cos(2 * M_PI * i / (height - 1))) *
                              0.5 * (1 - std::cos(2 * M_PI * j / (width - 1)));
    }

    features = features.cwiseProduct(_hann);

    return features;
}
